from Fcanet import *
from LoadData import *
import torch
import torch.nn.functional as F
import torch.optim as optim
from torchsummary import summary
from tqdm import tqdm
from Label_Smoothing import *
from Cutout import *

# cutout如何使用
cutout = Cutout(n_holes=1, length=8)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 顺滑标签如何使用
criterion = LabelSmoothingCrossEntropy()
# criterion = nn.CrossEntropyLoss

def train(net, trainloader, optimizer):
    net.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(tqdm(trainloader)):
        data, target = data.to(device), target.to(device)
        # 这里是使用cutout的地方
        # data = cutout(data)
        # data = data.to(device)
        optimizer.zero_grad()
        output = net(data)

        #loss = F.cross_entropy(output, target)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.data.item()

    running_loss /= len(train_loader)
    print('Train Loss: {:.6f}'.format(running_loss))


def test(net, testloader):
    net.eval()
    val_loss = 0
    correct = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(tqdm(testloader)):
            data, target = data.to(device), target.to(device)
            output = net(data)
            val_loss += criterion(output, target).data.item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()


    print("correct:", correct, ", total:", len(testloader.dataset))
    print('accuracy=', 100. * correct / len(testloader.dataset))
    print('val_loss=', val_loss/len(testloader))


if __name__ == '__main__':
    learning_rate = 1e-4  # 按照论文里设置的
    epoches = 20
    weight_decay = 5e-2


    # model = FcaResNet(FcaBottleNeck, [3, 4, 6, 3]).to(device)
    model = FcaResNet(FcaBasicBlock, [2, 2, 2, 2]).to(device)
    # model = FcaResNet(FcaBasicBlock, [3, 4, 6, 3]).to(device)
    # model = FcaRes2Net(Res2NetBottleneck, [3, 4, 6, 3]).to(device)
    # model = DenseNet(blocks=(6, 12, 32, 32), num_classes=100).to(device)
    # summary(model, input_size=(3, 32, 32))

    model.load_state_dict(torch.load('fcaresnet18cifar100.pkl'))  # 装载上传训练的参数
    model = model.to(device)


    # optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.5)
    # optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    optimizer_AdamW = torch.optim.AdamW(model.parameters(), weight_decay=weight_decay, lr=learning_rate)
    optimizer = optimizer_AdamW
    for epoch in range(epoches+1):
        print("epoch: ", epoch)
        train(model, train_loader, optimizer)
        torch.save(model.state_dict(), 'fcaresnet18cifar100.pkl')
        test(model, test_loader)

    #torch.save(model.state_dict(), 'fcaresnet18cifar100.pkl')